from pathlib import Path

import numpy as np
from docx.shared import Cm
from docxtpl import DocxTemplate, InlineImage
from matplotlib import pyplot as plt
from scipy import optimize
from scipy.special import jv  # Вводим функцию Бесселя

# Разработчики кода: Савчук и Кириенко

N = 5  # номер варианта
R0 = 0.1 * N  # радиус бака
fig_num = 1  # нумерация рисунков

x_bessel = np.arange(0.2, 10, 0.0001)

J1 = jv(1, x_bessel)


def paint_2d(xlabel: str, ylabel: str, x: np.ndarray, y: np.ndarray, fig_num: float):
    """Функция, строющая график функции

    Args:
        xlabel (str): наименование оси x
        ylabel (str): наименование оси y
        x (np.ndarray): вектор значений по x
        y (np.ndarray): вектор значений по y
        fig_num (float): номер рисунка

    Returns:
        float: увеличиваем номер рисунка на 1
    """
    plt.figure(fig_num)
    plt.grid(True)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.plot(x, y)
    fig_num += 1
    return fig_num


fig_num = paint_2d("x, м", "J1(x), б.р.", x_bessel, J1, fig_num)

# Определение корней функции Бесселя


def f_1(x_bessel: np.ndarray) -> np.ndarray:
    """Нахождение функции Бесселя при n=1

    Args:
        x_bessel (np.ndarray): координата, м

    Returns:
        np.ndarray: значения функции Бесселя, б. р.
    """
    return jv(1, x_bessel)


solution1 = optimize.root_scalar(f_1, bracket=[3, 4], method="bisect")
solution2 = optimize.root_scalar(f_1, bracket=[6, 8], method="bisect")
solution3 = optimize.root_scalar(f_1, bracket=[10, 12], method="bisect")
roots_1 = [solution1.root, solution2.root, solution3.root]


DJ1 = 1 / 2 * (jv(0, x_bessel) - jv(2, x_bessel))  # Производная функции Бесселя
fig_num = paint_2d("x, м", "DJ1(x), б.р.", x_bessel, DJ1, fig_num)

# Определение корней производной функции Бесселя


def df_1(x_bessel: np.ndarray) -> np.ndarray:
    """Нахождение производной функции Бесселя при n=1

    Args:
        x_bessel (np.ndarray): координата, м

    Returns:
        np.ndarray: значения производной функции Бесселя, б. р.
    """
    return 1 / 2 * (jv(0, x_bessel) - jv(2, x_bessel))


solution_ksi1 = optimize.root_scalar(df_1, bracket=[1, 2], method="bisect")
solution_ksi2 = optimize.root_scalar(df_1, bracket=[4, 6], method="bisect")
solution_ksi3 = optimize.root_scalar(df_1, bracket=[8, 10], method="bisect")
ksi = [solution_ksi1.root, solution_ksi2.root, solution_ksi3.root]


# Определение пси-функций


def psi(r: np.ndarray, teta: np.ndarray, ksi: float) -> np.ndarray:
    """Нахождение формы собственных колебаний свободной поверхности жидкости

    Args:
        r (np.ndarray): радиальная координата, м
        teta (np.ndarray): угловая координата, рад
        ksi (float): корни функции Бесселя для n=1, м

    Returns:
        np.ndarray: форма собственных колебаний свободной поверхности жидкости, б. р.
    """
    return jv(1, ksi * r / R0) / jv(1, ksi) * np.cos(teta)


# Создание массивов x, y, z для поверхностей вращения

r_array = np.arange(0, R0 + R0 / 250, R0 / 250)
teta_array = np.arange(0, 2 * np.pi + 2 * np.pi / 250, 2 * np.pi / 250)
r_array_2d = np.arange(-(R0), R0 + R0 / 250, R0 / 250)

x_graf_array = np.empty((len(r_array), len(teta_array)))
y_graf_array = np.empty((len(r_array), len(teta_array)))
z1_graf_array = np.empty((len(r_array), len(teta_array)))
z2_graf_array = np.empty((len(r_array), len(teta_array)))
z3_graf_array = np.empty((len(r_array), len(teta_array)))


for j in range(len(teta_array)):
    for i in range(len(r_array)):
        x_graf_array[i, j] = r_array[i] * np.cos(teta_array[j])
        y_graf_array[i, j] = r_array[i] * np.sin(teta_array[j])
        z1_graf_array[i, j] = psi(r_array[i], teta_array[j], ksi[0])
        z2_graf_array[i, j] = psi(r_array[i], teta_array[j], ksi[1])
        z3_graf_array[i, j] = psi(r_array[i], teta_array[j], ksi[2])


def paint_3d(zlabel: str, z: np.ndarray, fig_num: float):
    """Построение формы собственных колебаний свободной поверхности жидкости в
    пространстве

    Args:
        zlabel (str): наименование оси z
        z (np.ndarray): матрица формы колебаний
        fig_num (float): номер рисунка

    Returns:
        float: увеличиваем номер рисунка на 1
    """
    plt.figure(fig_num)
    ax = plt.axes(projection="3d")
    ax.set_xlabel("r, м")
    ax.set_ylabel("r, м")
    ax.set_zlabel(zlabel)
    ax.plot_surface(x_graf_array, y_graf_array, z)
    fig_num += 1
    return fig_num


# Построение графика и поверхности для пси 1
fig_num = paint_2d("r, м", "psi_11, м", r_array_2d, psi(r_array_2d, 0, ksi[0]), fig_num)
fig_num = paint_3d("psi_11, м", z1_graf_array, fig_num)

# Построение графика и поверхности для пси 2
fig_num = paint_2d("r, м", "psi_12, м", r_array_2d, psi(r_array_2d, 0, ksi[1]), fig_num)
fig_num = paint_3d("psi_12, м", z2_graf_array, fig_num)

# Построение графика и поверхности для пси 3
fig_num = paint_2d("r, м", "psi_13, м", r_array_2d, psi(r_array_2d, 0, ksi[2]), fig_num)
fig_num = paint_3d("psi_13, м", z3_graf_array, fig_num)


###################
# Для n-ого порядка#
###################

x_bessel = np.arange(0.2, 17, 0.0001)

JN = jv(N, x_bessel)
fig_num = paint_2d("x, м", "JN(x), б.р.", x_bessel, JN, fig_num)

# Определение корней функции Бесселя


def f_n(x_bessel: np.ndarray) -> np.ndarray:
    """Нахождение функции Бесселя при n=N

    Args:
        x_bessel (np.ndarray): координата, м

    Returns:
        np.ndarray: значения функции Бесселя, б. р.
    """

    return jv(N, x_bessel)


solution1 = optimize.root_scalar(f_n, bracket=[5, 9], method="bisect")
solution2 = optimize.root_scalar(f_n, bracket=[10, 13], method="bisect")
solution3 = optimize.root_scalar(f_n, bracket=[15, 17], method="bisect")
roots_n = [solution1.root, solution2.root, solution3.root]


DJN = 1 / 2 * (jv(N - 1, x_bessel) - jv(N + 1, x_bessel))  # Производная функции Бесселя
fig_num = paint_2d("x, м", "DJN(x), б.р.", x_bessel, DJN, fig_num)

# Определение корней производной функции Бесселя


def df_n(x_bessel: np.ndarray) -> np.ndarray:
    """Нахождение производной функции Бесселя при n=N

    Args:
        x_bessel (np.ndarray): координата, м

    Returns:
        np.ndarray: значения производной функции Бесселя, б. р.
    """

    return 1 / 2 * (jv(N - 1, x_bessel) - jv(N + 1, x_bessel))


solution_ksi1 = optimize.root_scalar(df_n, bracket=[5, 9], method="bisect")
solution_ksi2 = optimize.root_scalar(df_n, bracket=[10, 12], method="bisect")
solution_ksi3 = optimize.root_scalar(df_n, bracket=[13, 15], method="bisect")
ksi_n = [solution_ksi1.root, solution_ksi2.root, solution_ksi3.root]


# Определение пси-функций


def psi_n(r: np.ndarray, teta: np.ndarray, ksi_n: float) -> np.ndarray:
    """Нахождение формы собственных колебаний свободной поверхности жидкости

    Args:
        r (np.ndarray): радиальная координата, м
        teta (np.ndarray): угловая координата, рад
        ksi (float): корни функции Бесселя для n=N, м

    Returns:
        np.ndarray: форма собственных колебаний свободной поверхности жидкости, б. р.
    """

    return jv(N, ksi_n * r / R0) / jv(N, ksi_n) * np.cos(teta * N)


# Создание массивов z для поверхностей вращения

zn1_graf_array = np.empty((len(r_array), len(teta_array)))
zn2_graf_array = np.empty((len(r_array), len(teta_array)))
zn3_graf_array = np.empty((len(r_array), len(teta_array)))


for j in range(len(teta_array)):
    for i in range(len(r_array)):
        zn1_graf_array[i, j] = psi_n(r_array[i], teta_array[j], ksi_n[0])
        zn2_graf_array[i, j] = psi_n(r_array[i], teta_array[j], ksi_n[1])
        zn3_graf_array[i, j] = psi_n(r_array[i], teta_array[j], ksi_n[2])

# Построение графика и поверхности для пси n1
fig_num = paint_2d(
    "r, м", "psi_n1, м", r_array_2d, psi_n(r_array_2d, 0, ksi_n[0]), fig_num
)
fig_num = paint_3d("psi_n1, м", zn1_graf_array, fig_num)

# Построение графика и поверхности для пси n2
fig_num = paint_2d(
    "r, м", "psi_n2, м", r_array_2d, psi_n(r_array_2d, 0, ksi_n[1]), fig_num
)
fig_num = paint_3d("psi_n2, м", zn2_graf_array, fig_num)

# Построение графика и поверхности для пси n3
fig_num = paint_2d(
    "r, м", "psi_n3, м", r_array_2d, psi_n(r_array_2d, 0, ksi_n[2]), fig_num
)
fig_num = paint_3d("psi_n3, м", zn3_graf_array, fig_num)


# Заполнение шаблона

docx_template = DocxTemplate("dz/gau_9_sem/template.docx")


# Названия рисунков
image_names = [
    "im_bessel_func_1",
    "im_dif_bessel_func_1",
    "im_2d_psi_11",
    "im_3d_psi_11",
    "im_2d_psi_12",
    "im_3d_psi_12",
    "im_2d_psi_13",
    "im_3d_psi_13",
    "im_bessel_func_n",
    "im_dif_bessel_func_n",
    "im_2d_psi_n1",
    "im_3d_psi_n1",
    "im_2d_psi_n2",
    "im_3d_psi_n2",
    "im_2d_psi_n3",
    "im_3d_psi_n3",
]


# Сохранение графиков
for i in range(16):
    fig = plt.figure(i + 1)
    fig.savefig(Path("dz/gau_9_sem/" + image_names[i] + ".png"), dpi=fig.dpi)

images = {
    image_name: InlineImage(
        docx_template, "dz/gau_9_sem/" + image_name + ".png", Cm(12)
    )
    for image_name in image_names
}

context = {
    "radius_baka": R0,
    "nomer_varianta": N,
    "korni_bessel_func_11": roots_1[0],
    "korni_bessel_func_12": roots_1[1],
    "korni_bessel_func_13": roots_1[2],
    "korni_bessel_func_n1": roots_n[0],
    "korni_bessel_func_n2": roots_n[1],
    "korni_bessel_func_n3": roots_n[2],
    "ksi_11": ksi[0],
    "ksi_12": ksi[1],
    "ksi_13": ksi[2],
    "ksi_n1": ksi_n[0],
    "ksi_n2": ksi_n[1],
    "ksi_n3": ksi_n[2],
}

context = context | images


docx_template.render(context)
docx_template.save("dz/gau_9_sem/Otchyot.docx")
